package com.ask.forMe.langchain4j.service;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.HuggingFaceTokenizer;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
public class EmbeddingService {
    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private EmbeddingStore embeddingStore;

    @Autowired
    private EmbeddingStore quaStore;

    /**
     * 测试向量模型
     */
    public void testEmbeddingModel() {
        Response<Embedding> embed = embeddingModel.embed("你好");
        System.out.println("向量维度：" + embed.content().vector().length);
        System.out.println("向量输出：" + embed.toString());
    }

    /**
     * 测试Pinecone向量数据库存储
     */
    public void testPineconeEmbeded() {
        TextSegment segment1 = TextSegment.from("我们的优化政策是：xxxx");
        Embedding embedding1 = embeddingModel.embed(segment1).content();
        quaStore.add(embedding1, segment1);

        TextSegment segment2 = TextSegment.from("我们的调度策略是：xxx");
        Embedding embedding2 = embeddingModel.embed(segment2).content();
        quaStore.add(embedding2, segment2);
    }


    /**
     * 相似度匹配
     */
    public EmbeddingMatch<TextSegment> embeddingSearch(String message) {

        // 1. 将查询问题转换为向量表示
        Embedding queryEmbedding = embeddingModel.embed(message).content();
        //  2. 构建搜索请求
        EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)  // 设置查询向量
                .maxResults(1)  // 只返回最相似的1个结果
                //.minScore(0.8)   // 相似度阈值
                .build();

        // 3. 执行向量相似度搜索
        EmbeddingSearchResult<TextSegment> searchResult = quaStore.search(searchRequest);

        // 4. 获取最匹配的结果
        EmbeddingMatch<TextSegment> embeddingMatch = searchResult.matches().get(0);

        // 5. 输出相似度分数
        System.out.println(embeddingMatch.score());

        // 6. 输出匹配的文本内容
        System.out.println(embeddingMatch.embedded().text());

        return embeddingMatch;
    }

    /**
     * 上传知识库到Pinecone
     */
    public void uploadKnowledge2Pinecone(String path, List<String> fileNames) {
        List<Document> documentList = fileNames
                .stream()
                .map(name -> FileSystemDocumentLoader.loadDocument(path + name))
                .collect(Collectors.toList());

        DocumentByParagraphSplitter documentSplitter = new DocumentByParagraphSplitter(
                300,
                30,
                new HuggingFaceTokenizer()
        );


        EmbeddingStoreIngestor.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .documentSplitter(documentSplitter)
                .build()
                .ingest(documentList);
    }

    public String querySimilarQuestions(String question) {

        // 1. 将查询问题转换为向量表示
        Embedding queryEmbedding = embeddingModel.embed(question).content();
        //  2. 构建搜索请求
        EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)  // 设置查询向量
                .maxResults(1)  // 只返回最相似的1个结果
                .minScore(0.8)   // 相似度阈值
                .build();

        // 3. 执行向量相似度搜索
        EmbeddingSearchResult<TextSegment> searchResult = quaStore.search(searchRequest);

        // 没有查找到相似向量
        if (searchResult.matches().isEmpty()) {
            return null;
        }
        // 4. 获取最匹配的结果
        EmbeddingMatch<TextSegment> embeddingMatch = searchResult.matches().get(0);

        // 5. 输出相似度分数
        System.out.println(embeddingMatch.score());

        // 6. 输出匹配的文本内容
        System.out.println(embeddingMatch.embedded().text());

        return embeddingMatch.embedded().text();
    }

    /**
     * 上传问题答案至向量数据库Pinecone
     *
     * @param ans
     */
    public void uploadAns2Pinecone(String ans) {
        TextSegment segment = TextSegment.from(ans);
        Embedding embedding = embeddingModel.embed(segment).content();
        quaStore.add(embedding, segment);
    }

    public void uploadAndFilter() {
        // 1. 创建带有元数据的文档集合
        List<Document> documents = Arrays.asList(
                createDocument("Java性能优化指南", "这是关于Java应用性能优化的详细指南...",
                        "技术文档", "Java", "2023", "已发布"),

                createDocument("Spring Boot最佳实践", "本文介绍Spring Boot应用开发的最佳实践...",
                        "技术文档", "Spring", "2023", "已发布"),

                createDocument("数据库设计草案", "初步的数据库结构设计和ER图...",
                        "设计文档", "数据库", "2023", "草稿"),

                createDocument("系统架构设计", "系统整体架构设计和组件交互关系...",
                        "架构文档", "系统设计", "2022", "已发布")
        );

        Document document = documents.get(0);
        System.out.println(document.metadata().getString("title"));
        //System.out.println(documents);
    }

    private static Document createDocument(String title, String content,
                                           String docType, String topic,
                                           String year, String status) {

        Map<String, String> metadataMap = new HashMap<>();
        metadataMap.put("title", title);
        metadataMap.put("docType", docType);
        metadataMap.put("topic", topic);
        metadataMap.put("year", year);
        metadataMap.put("status", status);

        Metadata metadata = Metadata.from(metadataMap);

        return Document.from(content, metadata);
    }
}
